// *************************************************************************************************
//
// MK3D Engine
// --------------------------------------
// Copyright (C) 2007-2011 Zeb (mailto: zebbey@gmail.com)
//
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//
// *************************************************************************************************

#include "mk_header.h"
#include "mk_pp_bloom.h"
#include "mk_render_device.h"
#include "mk_engine.h"
#include "mk_error.h"
#include "mk_shader.h"
#include "mk_render_window.h"
#include "mk_resource_manager.h"
#include "mk_painter.h"
#include "mk_math.h"

namespace MK3DE
{
	CPostProcessorBloom::CPostProcessorBloom()
	{
		_brightPassRT = NULL;
		memset(_bloomRT, 0, sizeof(_bloomRT));
		memset(_brightPassSampleOffsets, 0, sizeof(_brightPassSampleOffsets));
		memset(_bloomSampleOffsetsV, 0, sizeof(_bloomSampleOffsetsV));
		memset(_bloomSampleWeightsV, 0, sizeof(_bloomSampleWeightsV));
		memset(_bloomSampleOffsetsH, 0, sizeof(_bloomSampleOffsetsH));
		memset(_bloomSampleWeightsH, 0, sizeof(_bloomSampleWeightsH));
		_shader = NULL;
		_vbQuad2D = NULL;
		_renderInput = NULL;
	}

	CPostProcessorBloom::~CPostProcessorBloom()
	{
		CRenderDevice* renderDevice = GetEngine()->GetRenderDevice();
		renderDevice->DestroyRenderTarget(_brightPassRT);
		renderDevice->DestroyRenderTarget(_bloomRT[0]);
		renderDevice->DestroyRenderTarget(_bloomRT[1]);
		SAFE_RELEASE(_shader);
		SAFE_RELEASE(_vbQuad2D);
		renderDevice->DestroyRenderInput(_renderInput);
	}

	BOOL CPostProcessorBloom::ApplyEffect(CRenderTarget* source, CRenderTarget* dest)
	{
		if (!_brightPassRT || !_bloomRT[0] || !_bloomRT[1])
		{
			if (!__Create())
				return FALSE;
		}

		BOOL ret = FALSE;

		ret = __BrightPass(source);
		MK_PROCESS_ERROR(ret);

		ret = __BloomPass();
		MK_PROCESS_ERROR(ret);

		ret = __FinalPass(source, dest);
		MK_PROCESS_ERROR(ret);

		ret = TRUE;
Exit0:
		return ret;
	}

	BOOL CPostProcessorBloom::__FinalPass(CRenderTarget* source, CRenderTarget* dest)
	{
		CRenderDevice* renderDevice = GetEngine()->GetRenderDevice();
		CTexture2D* sourceTexture = source->GetRenderTargetTexture();
		float width = (float)sourceTexture->GetWidth();
		float height = (float)sourceTexture->GetHeight();

		DWORD passNum;
		BOOL ret = renderDevice->SetRenderTarget(dest);
		MK_PROCESS_ERROR(ret);

		ret = _shader->SetTexture("sourceTexture", sourceTexture);
		MK_PROCESS_ERROR(ret);

		ret = _shader->SetTexture("bloomTexture", _bloomRT[0]->GetRenderTargetTexture());
		MK_PROCESS_ERROR(ret);

		ret = _shader->TechniqueBegin("FinalPass", passNum);
		MK_PROCESS_ERROR(ret);

		for (DWORD pass = 0; pass < passNum; ++pass)
		{
			if (_shader->PassBegin(pass))
			{
				__DrawQuad(width, height);

				_shader->PassEnd();
			}
		}

		_shader->TechniqueEnd();

		ret = TRUE;
Exit0:
		return ret;
	}

	BOOL CPostProcessorBloom::__BrightPass(CRenderTarget* source)
	{
		CRenderDevice* renderDevice = GetEngine()->GetRenderDevice();
		CTexture2D* renderTargetTexture = _brightPassRT->GetRenderTargetTexture();
		float width = (float)renderTargetTexture->GetWidth();
		float height = (float)renderTargetTexture->GetHeight();

		DWORD passNum;
		BOOL ret = renderDevice->SetRenderTarget(_brightPassRT);
		MK_PROCESS_ERROR(ret);

		ret = _shader->SetFloatArray("sampleOffsets", (float*)_brightPassSampleOffsets, 32);
		MK_PROCESS_ERROR(ret);

		ret = _shader->SetTexture("sourceTexture", source->GetRenderTargetTexture());
		MK_PROCESS_ERROR(ret);

		ret = _shader->TechniqueBegin("DownScale3X3_BrightPass", passNum);
		MK_PROCESS_ERROR(ret);

		for (DWORD pass = 0; pass < passNum; ++pass)
		{
			if (_shader->PassBegin(pass))
			{
				__DrawQuad(width, height);

				_shader->PassEnd();
			}
		}

		_shader->TechniqueEnd();

		ret = TRUE;
Exit0:
		return ret;
	}

	BOOL CPostProcessorBloom::__DrawQuad(float width, float height)
	{
		BOOL ret = FALSE;
		CRenderDevice* renderDevice = GetEngine()->GetRenderDevice();
		VERT_2D_TEX1 vert[4];
		vert[0].pos = D3DXVECTOR4(-0.5f, -0.5f, 0, 1.0f);
		vert[1].pos = D3DXVECTOR4(width - 0.5f, -0.5f, 0, 1.0f);
		vert[2].pos = D3DXVECTOR4(-0.5f, height - 0.5f, 0, 1.0f);
		vert[3].pos = D3DXVECTOR4(width - 0.5f, height - 0.5f, 0, 1.0f);

		vert[0].uv = D3DXVECTOR2(0, 0);
		vert[1].uv = D3DXVECTOR2(1.0f, 0);
		vert[2].uv = D3DXVECTOR2(0, 1.0f);
		vert[3].uv = D3DXVECTOR2(1.0f, 1.0f);

		VERT_2D_TEX1* p;
		_vbQuad2D->Lock((void*&)p, 0, 0, LO_DISCARD);
		memcpy(p, vert, sizeof(VERT_2D_TEX1) * 4);
		_vbQuad2D->Unlock();

		ret = renderDevice->BindRenderInput(_renderInput);
		MK_PROCESS_ERROR(ret);

		ret = renderDevice->DrawPrimitive(PT_TRIANGLESTRIP, 0, 0, 4, 0, 2);
		MK_PROCESS_ERROR(ret);

		ret = TRUE;
Exit0:
		return ret;
	}

	BOOL CPostProcessorBloom::__BloomPass()
	{
		CRenderDevice* renderDevice = GetEngine()->GetRenderDevice();
		CTexture2D* bloomTexture = _bloomRT[1]->GetRenderTargetTexture();
		float width = (float)bloomTexture->GetWidth();
		float height = (float)bloomTexture->GetHeight();

		// Horizontal blur
		DWORD passNum;
		BOOL ret = renderDevice->SetRenderTarget(_bloomRT[1]);
		MK_PROCESS_ERROR(ret);

		ret = _shader->SetFloatArray("sampleOffsets", (float*)_bloomSampleOffsetsH, 32);
		MK_PROCESS_ERROR(ret);

		ret = _shader->SetFloatArray("sampleWeights", (float*)_bloomSampleWeightsH, 64);

		ret = _shader->SetTexture("sourceTexture", _brightPassRT->GetRenderTargetTexture());
		MK_PROCESS_ERROR(ret);

		ret = _shader->TechniqueBegin("Bloom", passNum);
		MK_PROCESS_ERROR(ret);

		for (DWORD pass = 0; pass < passNum; ++pass)
		{
			if (_shader->PassBegin(pass))
			{
				__DrawQuad(width, height);

				_shader->PassEnd();
			}
		}

		_shader->TechniqueEnd();

		// Vertical blur
		bloomTexture = _bloomRT[0]->GetRenderTargetTexture();
		width = (float)bloomTexture->GetWidth();
		height = (float)bloomTexture->GetHeight();

		ret = renderDevice->SetRenderTarget(_bloomRT[0]);
		MK_PROCESS_ERROR(ret);

		ret = _shader->SetFloatArray("sampleOffsets", (float*)_bloomSampleOffsetsV, 32);
		MK_PROCESS_ERROR(ret);

		ret = _shader->SetFloatArray("sampleWeights", (float*)_bloomSampleWeightsV, 64);

		ret = _shader->SetTexture("sourceTexture", _bloomRT[1]->GetRenderTargetTexture());
		MK_PROCESS_ERROR(ret);

		ret = _shader->TechniqueBegin("Bloom", passNum);
		MK_PROCESS_ERROR(ret);

		for (DWORD pass = 0; pass < passNum; ++pass)
		{
			if (_shader->PassBegin(pass))
			{
				__DrawQuad(width, height);

				_shader->PassEnd();
			}
		}

		_shader->TechniqueEnd();

		ret = TRUE;
Exit0:
		return ret;
	}

	BOOL CPostProcessorBloom::OnDeviceLost()
	{
		CRenderDevice* renderDevice = GetEngine()->GetRenderDevice();
		renderDevice->DestroyRenderTarget(_brightPassRT);
		renderDevice->DestroyRenderTarget(_bloomRT[0]);
		renderDevice->DestroyRenderTarget(_bloomRT[1]);
		return TRUE;
	}

	BOOL CPostProcessorBloom::__Create()
	{
		BOOL ret = FALSE;
		CEngine* engine = GetEngine();
		CRenderDevice* renderDevice = engine->GetRenderDevice();
		CResourceManager* resourceManager = engine->GetResourceManager();
		CRenderWindow* renderWindow = engine->GetCurrentRenderWindow();
		DWORD texWidth = renderWindow->GetWidth() / 8;
		DWORD texHeight = renderWindow->GetHeight() / 8;

		renderDevice->DestroyRenderTarget(_brightPassRT);
		renderDevice->DestroyRenderTarget(_bloomRT[0]);
		renderDevice->DestroyRenderTarget(_bloomRT[1]);

		// Create all render target
		ret = renderDevice->CreateRenderTarget(texWidth, texHeight, PF_A8R8G8B8, FALSE, _brightPassRT);
		MK_PROCESS_ERROR(ret);

		ret = renderDevice->CreateRenderTarget(texWidth, texHeight, PF_A8R8G8B8, FALSE, _bloomRT[0]);
		MK_PROCESS_ERROR(ret);

		ret = renderDevice->CreateRenderTarget(texWidth, texHeight, PF_A8R8G8B8, FALSE, _bloomRT[1]);
		MK_PROCESS_ERROR(ret);

		if (!_shader)
		{
			// Load shader file.
			ret = resourceManager->GetShaderFromFile("shader/bloom.fx", _shader);
			MK_PROCESS_ERROR(ret);
		}

		if (!_vbQuad2D)
		{
			// Create vertex buffer and render input.
			ret = resourceManager->GetVertexBuffer(sizeof(VERT_2D_TEX1), 4, (HARDWARE_BUFFER_USAGE)(HBU_DYNAMIC|HBU_WRITE_ONLY), _vbQuad2D);
			MK_PROCESS_ERROR(ret);
		}

		if (!_renderInput)
		{
			RenderInputElemDesc rieQuad2D[] = { 
				RenderInputElemDesc(SU_POSITION_T, 0, 0, 0),
				RenderInputElemDesc(SU_TEXCOORD, 0, 16, 0)
			};

			ret = renderDevice->CreateRenderInput(rieQuad2D, 2, _renderInput);
			MK_PROCESS_ERROR(ret);

			ret = _renderInput->SetVertexBuffer(0, _vbQuad2D);
			MK_PROCESS_ERROR(ret);
		}

		// Calculate bright pass down scale sample offset.
		__GetSampleOffsetsDownScale3x3(renderWindow->GetWidth() / 2, renderWindow->GetHeight() / 2);

		__GetSampleOffsets_Bloom(texWidth, _bloomSampleOffsetsH, _bloomSampleWeightsH, 3.0f, 1.25f, FALSE);
		__GetSampleOffsets_Bloom(texHeight, _bloomSampleOffsetsV,  _bloomSampleWeightsV, 3.0f, 1.25f, TRUE);

		ret = TRUE;
Exit0:
		return ret;
	}

	void CPostProcessorBloom::__GetSampleOffsets_Bloom(DWORD dwD3DTexSize, D3DXVECTOR2* avSampleOffsets, D3DXVECTOR4* avColorWeight, float fDeviation, float fMultiplier, BOOL bVertical)
	{
		float afTexCoordOffset[15];
		int i = 0;
		float tu = 1.0f / (float)dwD3DTexSize;

		// Fill the center texel
		float weight = 1.0f * GaussianDistribution2(0, 0, fDeviation);
		avColorWeight[0] = D3DXVECTOR4(weight, weight, weight, 1.0f);

		afTexCoordOffset[0] = 0.0f;

		// Fill the right side
		for (i = 1; i < 8; i++)
		{
			weight = fMultiplier * GaussianDistribution2((float)i, 0, fDeviation);
			afTexCoordOffset[i] = i * tu;

			avColorWeight[i] = D3DXVECTOR4(weight, weight, weight, 1.0f);
		}

		// Copy to the left side
		for( i = 8; i < 15; i++ )
		{
			avColorWeight[i] = avColorWeight[i - 7];
			afTexCoordOffset[i] = -afTexCoordOffset[i - 7];
		}

		if (bVertical)
		{
			for( i = 0; i < 16; i++ )
				avSampleOffsets[i] = D3DXVECTOR2(0.0f, afTexCoordOffset[i]);
		}
		else
		{
			for( i = 0; i < 16; i++ )
				avSampleOffsets[i] = D3DXVECTOR2(afTexCoordOffset[i], 0.0f);
		}
	}

	void CPostProcessorBloom::__GetSampleOffsetsDownScale3x3(DWORD width, DWORD height)
	{
		float tU = 1.0f / width;
		float tV = 1.0f / height;

		// Sample from the 9 surrounding points. 
		int index = 0;
		for (int y = -1; y <= 1; y++)
		{
			for (int x = -1; x <= 1; x++)
			{
				_brightPassSampleOffsets[index].x = x * tU;
				_brightPassSampleOffsets[index].y = y * tV;

				index++;
			}
		}
	}
}